local RBTnode = {}

local function RBTNode(key, val)
    local o = {}
    setmetatable(o,RBTnode)
    RBTnode.__index = RBTnode
    o.left = {}
    o.right = {}
    o.parent = {}
    o.key = key or {}
    o.value = val or {}
    o.color = {}    --Boolean 类型  black：true red：false
    return o
end

local RBtree = {}

--左旋
function RBtree:leftRotate(node)
    if not self:find(node.key) then
        error("左旋节点不存在")
    end
    local temp_node = node.right
    node.right = temp_node.left
    if temp_node.left.key ~= nil then
        temp_node.left.parent = node
    end
    temp_node.parent = node.parent
    if node.parent.key == nil then
        self.rootNode = temp_node
    else
        if node.parent.left.key == node.key then
            node.parent.left = temp_node
        else
            node.parent.right = temp_node
        end
    end
    temp_node.left = node
    node.parent = temp_node
    -- print("转完之后的key值"..temp_node.key)
end

-- 右旋
function RBtree:rightRotate(node)
    if not self:find(node.key) then
        error("右旋节点不存在")
    end
    local temp_node = node.left
    node.left = temp_node.right
    if temp_node.right.key ~= nil then
        temp_node.right.parent = node
    end
    temp_node.parent = node.parent
    if node.parent.key == nil then
        self.rootNode = temp_node
    else
        if node.key == node.parent.right.key then
            node.parent.right = temp_node
        else
            node.parent.left = temp_node
        end
    end
    temp_node.right = node
    node.parent  = temp_node
end

--添加函数
function RBtree:insert(key, value)
    --判断子节点的内
    if type(key) ~= "number" or value == nil then
        error("请规范输入参数 key：number  value：~=nil")
    end
    if self:find(key) then
        print("插入了相同的key: "..key)
        return
    end
    local rootNode = self.rootNode
    --当根节点为空时
    if not next(rootNode) then
        self.rootNode = RBTNode(key, value)
        return
    end
    local temp_node = {}
    local node = RBTNode(key, value)
    -- 插入到树中
    while rootNode.right ~= nil do
        temp_node = rootNode
        if node.key < rootNode.key then
            rootNode = rootNode.left
        else
            rootNode = rootNode.right
        end
    end
    node.parent = temp_node
    if temp_node ~= nil then
        if node.key < temp_node.key then
            temp_node.left = node
        else
            temp_node.right = node
        end
    else
        self.rootNode = node
    end
    --节点颜色改为红色
    node.color = false
    --修正为红黑树
    self:insertFixUp(node)
end

--添加之后的修正函数
function RBtree:insertFixUp(node)
    local parent = node.parent
    local gparent = {}
    while parent.key ~= nil and not parent.color do
        gparent = parent.parent
        if parent.key == gparent.left.key then
            local uncle = gparent.right
            --叔叔的节点是红色时
            if uncle.key ~= nil and not uncle.color then
                uncle.color = true
                parent.color = true
                gparent.color = false
                node = gparent
                goto continue
            end
            --叔叔是黑色节点，传入的节点为右孩子时 左旋
            if parent.right.key == node.key then
                self:leftRotate(parent)
                local temp = parent
                parent = node
                node  = temp
            end
            --叔叔是黑色节点，传入的节点为左孩子时 右旋
            parent.color = true
            gparent.color = false
            self:rightRotate(gparent)
        else
            local uncle = gparent.left
            --叔叔节点为红色时
            if uncle.key ~= nil and not uncle.color then
                uncle.color = true
                parent.color = true
                gparent.color = false
                node = gparent
                goto continue
            end
            --叔叔是黑色节点，传入的节点为左孩子时 右旋
            if parent.left.key == node.key then
                self:rightRotate(parent)
                local temp = parent
                parent = node
                node = temp
            end
            --叔叔是黑色节点，传入的节点为左孩子时 左旋
            parent.color = true
            gparent.color = false
            self:leftRotate(gparent)
        end
        ::continue::
        parent = node.parent
    end
    self.rootNode.color = true
end

--移除函数
function RBtree:remove(key)
    if key == nil or type(key) ~= "number" then
        print("请正确输入要移除的键值")
        return
    end
    local node = self:_find(key)
    if node == nil then
        print("未找到要移除的键值对")
        return
    end
    local child = {}
    local parent = {}
    local is_fix = false
    --左右节点不为空时
    if node.left.key ~= nil and node.right.key ~= nil then
        --取代节点
        local replace = node
        replace = replace.right
        while replace.left.key ~= nil do
            replace = replace.left
        end
        --根节点不存在父亲节点
        if node.parent.key ~= nil then
            if node.parent.left.key == node.key then
                node.parent.left = replace
            else 
                node.parent.right = replace
            end
        else
            self.rootNode = replace
        end
        child = replace.right
        parent = replace.parent
        is_fix = replace.color
        --删除的节点是后继节点的父亲节点
        if parent.key == node.key then
            parent = replace
        else 
            if child.key ~= nil then
                child.parent = parent
            end
            parent.left = child
            replace.right = node.right
            node.right.parent = replace
        end
        replace.parent = node.parent
        replace.color = node.color
        replace.left = node.left
        node.left.parent = replace
        if is_fix then 
            self:removeFixUp(child, parent)
        end
        node = {}
        return
    end
    if node.left.key ~= nil then
        child = node.left
    else 
        child = node.right
    end
    parent = node.parent
    is_fix = node.color
    if child.key ~= nil then
        child.parent = parent
    end
    if parent.key ~= nil then
        if parent.left.key == node.key then
            parent.left = child
        else 
            parent.right = child
        end
    else
        self.rootNode = child
    end
    if is_fix then
        self:removeFixUp(child,parent)
    end
    node = {}
end

--移除修正函数
function RBtree:removeFixUp(node, parent)
    local temp_node = {}
    while (node.key == nil or node.color) and node.key ~= self.rootNode.key do
        if parent.left.key == node.key then
            temp_node = parent.right
            --兄弟节点为红色时
            if not temp_node.color then
                temp_node.color = true
                parent.color = false
                self:leftRotate(parent)
                temp_node = parent.right
            end
            --兄弟是黑色，两个孩子也是黑色
            if ((temp_node.left.key == nil or temp_node.left.color) and (temp_node.right.key == nil or temp_node.color)) then
                temp_node.color = false
                node = parent
                parent = node.parent
            else
                --兄弟节点黑色，兄弟的左孩子红色右孩子黑色
                if temp_node.right.key == nil or temp_node.color then
                    temp_node.left.color = true
                    temp_node.color = false
                    self:rightRotate(temp_node)
                    temp_node = parent.right
                end
                --兄弟节点黑色，兄弟的右孩子黑色
                temp_node.color = parent.color
                parent.color = true
                temp_node.right.color = true
                self:leftRotate(parent)
                node = self.rootNode
                break
            end
        else
            temp_node = parent.left
            --兄弟节点为红色
            if not temp_node.color then
                temp_node.color = true
                parent.color = false
                self:rightRotate(parent)
                temp_node = parent.left
            end
            --兄弟节点黑色两个孩子黑色
            if ((temp_node.left.key ==nil or temp_node.color) and (temp_node.right.key == nil or temp_node.right.color)) then
                temp_node.color = false
                node = parent
                parent = node.parent
            else 
                --兄弟节点黑色，兄弟的左孩子红色右孩子黑色
                if temp_node.left.key == nil or temp_node.left.color then
                    temp_node.right.color = true
                    temp_node.color = false
                    self:leftRotate(temp_node)
                    temp_node = parent.left
                end
                --兄弟节点黑色，兄弟的右孩子黑色
                temp_node.color = parent.color
                parent.color = true
                temp_node.left.color = true
                self:rightRotate(parent)
                node = self.rootNode
                break
            end
        end
    end
    if node.key ~= nil then
        node.color = true
    end
end

--查找函数返回节点
function RBtree:_find(key)
    if key == nil or type(key) ~= "number" then
        error("请正确输入key值")
    end
    function Search(node, key)
        if node.key == nil then
            print("节点未在RBtree中")
            return {}
        end
        if key < node.key then
            return Search(node.left, key)
        elseif key > node.key then
            return Search(node.right, key)
        else
            return node
        end
    end
    return Search(self.rootNode, key)
end

--查找函数返回Boolean
function RBtree:find(key)
    -- print(key)
    if key == nil or type(key) ~= "number" then
        error("请正确输入key值")
    end
    function Search(node, key)
        if node.key == nil then
            return false
        end
        if key < node.key then
            return Search(node.left, key)
        elseif key > node.key then
            return Search(node.right, key)
        else
            return true
        end
    end
    return Search(self.rootNode, key)
end


function RBtree:clear()
  self.rootNode= {}
end

--递归遍历中序输出
function RBtree:print()
    local function inOrder(node)
        if node.left ~= nil then
            inOrder(node.left)
        end
        if node.value ~= nil then
            if self.rootNode.key == node.key then
                print("key: "..node.key.." value: "..node.value.."  头结点")
            else 
                print("key: "..node.key.." value: "..node.value)
            end
        end
        if node.right ~= nil then
            inOrder(node.right)
        end
    end
    inOrder(self.rootNode)
    print("==============分界线=============")
end

local function RBTree()
    local o ={}
    setmetatable(o, RBtree)
    RBtree.__index = RBtree
    o.rootNode = {}
    return o
end

return RBTree


